import base64
import json
from lambda_helper.s3_boto import UserInfoDynamoDB, is_running_on_aws_lambda
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
if not is_running_on_aws_lambda():
    handler = logging.StreamHandler()
    # format the log
    formatter = logging.Formatter('%(name)s - - [%(asctime)s] "%(message)s" %(levelname)%')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

def extract_api_key(headers):
    """从请求头中提取API Key"""
    api_key = None
    for key in ['apikey', 'ApiKey', 'Apikey']:
        if key in headers:
            api_key = headers[key]
            break
    logger.info(f"API key extracted: {api_key[:4]}...{api_key[-4:]}" if api_key else "No API key found")
    return api_key

def decode_jwt_payload(token):
    """解码JWT token的有效载荷部分"""
    try:
        # 分离JWT的三个部分
        _, payload_encoded, _ = token.split('.')
        # 添加URL安全的base64编码所需的填充字符
        payload_encoded += '=' * (4 - len(payload_encoded) % 4)
        # 处理URL安全字符
        payload_encoded = payload_encoded.replace('-', '+').replace('_', '/')
        # 解码并解析JSON
        decoded_payload = base64.b64decode(payload_encoded).decode('utf-8')
        return json.loads(decoded_payload)
    except Exception as e:
        logger.error(f"JWT decoding failed: {str(e)}")
        return None


class TokenData:

    ZhuChunDaShuaige = 'zhuchundashuaige'
    Msal = 'mirosoft'

    def __init__(self, api_key):
        self.api_key = api_key
        self.access_token = None
        self.platfrom = None
        self.user_id = None
        self.platfrom_id = None

    def parse_api_key(self):
        try:
            user_id, access_token, platform, platform_user_id = base64.b64decode(self.api_key).decode('utf-8').split('$$')
            self.user_id = user_id
            self.access_token = access_token
            self.platfrom = platform
            self.platfrom_id = platform_user_id
        except Exception as e:
            logger.error(f"API key parsing failed: {str(e)}")

    def is_valid(self):
        no_empty = self.user_id and self.access_token and self.platfrom and self.platfrom_id
        if not no_empty:
            return False
        if self.platfrom not in [self.ZhuChunDaShuaige, self.Msal]:
            return False
        user_info = UserInfoDynamoDB().get(self.user_id)
        self.user_info = user_info
        if not user_info:
            return False
        return True


def extract_data_from_header(headers: dict):
    api_key = extract_api_key(headers)
    token_data = TokenData(api_key)
    if not api_key:
        return token_data
    return None